package com.base.service;

import com.base.bean.EsException;
import com.base.bean.EsQueryReqPO;
import com.base.bean.EsQueryRespPO;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.springframework.stereotype.Repository;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * es查询持久层
 * https://blog.csdn.net/yzh_1346983557/article/details/122674983
 * @author yangzihe
 * @date 2022/3/18
 */
@Repository
@Slf4j
public class EsSearchRepository {

    @Resource
    private RestHighLevelClient highLevelClient;

    /**
     * 统计数量
     *
     * @param index 索引数组
     * @param query 查询条件
     *
     * @return 数量
     */
    public Long count(String[] index, QueryBuilder query) {
        // 请求对象
        CountRequest countRequest = new CountRequest(index, query);

        // 查询结果
        CountResponse countResponse;
        try {
            log.info("es统计数量请求：index={}, query=\n{}", index, query);
            countResponse = highLevelClient.count(countRequest, RequestOptions.DEFAULT);
            log.info("es统计数量结果：{}", countResponse);
        } catch (IOException e) {
            log.error("es统计数量，IO异常！index={}, query={}", index, query, e);
            throw new EsException("es统计数量，IO异常！");
        }

        if (RestStatus.OK.equals(countResponse.status())) {
            return countResponse.getCount();
        } else {
            log.error("es查询返回的状态码异常！status={}, index={}, query={}", countResponse.status(), index, query);
            throw new EsException("es统计数量返回的状态码异常");
        }
    }

    /**
     * 查询
     *
     * @param queryPO 查询条件
     *
     * @return 查询结果Map
     */
    public List<Map<String, Object>> search(EsQueryReqPO queryPO) {
        // 封装查询源对象
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 查询条件
        sourceBuilder.query(queryPO.getQuery());

        // 排序字段
        if (StringUtils.isNotBlank(queryPO.getSortField()) && queryPO.getSort() != null) {
            FieldSortBuilder order = new FieldSortBuilder(queryPO.getSortField()).order(queryPO.getSort());
            sourceBuilder.sort(order);
        }

        // 页大小，默认返回10条
        if (queryPO.getPageSize() != null) {
            sourceBuilder.size(queryPO.getPageSize());
        }

        // 设置索引、source
        SearchRequest searchRequest = new SearchRequest(queryPO.getIndex()).source(sourceBuilder);

        // 查询结果
        SearchResponse searchResponse;
        try {
            log.info("es查询请求：index={}, 请求source={}", queryPO.getIndex(), searchRequest.source());
            searchResponse = highLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info("es查询结果：{}", searchResponse);
        } catch (IOException e) {
            log.error("es查询，IO异常！es查询index={}, 请求source={}", queryPO.getIndex(), searchRequest.source(), e);
            throw new EsException("es查询，IO异常！");
        }

        if (RestStatus.OK.equals(searchResponse.status())) {
            // 解析对象
            SearchHit[] hits = searchResponse.getHits().getHits();
            // 获取source
            return Arrays.stream(hits).map(SearchHit::getSourceAsMap).collect(Collectors.toList());
        } else {
            log.error("es查询返回的状态码异常！searchResponse.status={}, index={}, 请求source={}", searchResponse.status(),
                    queryPO.getIndex(), searchRequest.source());
            throw new EsException("es查询返回的状态码异常");
        }
    }

    /**
     * 分页查询
     *
     * @param queryPO 分页查询对象
     *
     * @return 分页查询结果
     */
    public EsQueryRespPO pageSearch(EsQueryReqPO queryPO) {
        // 默认分页参数设置
        if (queryPO.getPageNum() == null) {
            queryPO.setPageNum(1);
        }
        if (queryPO.getPageSize() == null) {
            queryPO.setPageSize(10);
        }

        // 封装查询源对象
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();

        // 查询条件
        sourceBuilder.query(queryPO.getQuery());

        // 排序字段
        if (StringUtils.isNotBlank(queryPO.getSortField()) && queryPO.getSort() != null) {
            FieldSortBuilder order = new FieldSortBuilder(queryPO.getSortField()).order(queryPO.getSort());
            sourceBuilder.sort(order);
        }

        // 开始行数，默认0
        sourceBuilder.from((queryPO.getPageNum() - 1) * queryPO.getPageSize());
        // 页大小，默认10
        sourceBuilder.size(queryPO.getPageSize());

        // 设置索引、source
        SearchRequest searchRequest = new SearchRequest(queryPO.getIndex()).source(sourceBuilder);

        // 查询结果
        SearchResponse searchResponse;
        try {
            log.info("es分页查询请求：index={}, 请求source={}", queryPO.getIndex(), searchRequest.source());
            searchResponse = highLevelClient.search(searchRequest, RequestOptions.DEFAULT);
            log.info("es分页查询结果：{}", searchResponse);
        } catch (IOException e) {
            log.error("es查询，IO异常！es查询index={}, 请求source={}", queryPO.getIndex(), searchRequest.source(), e);
            throw new EsException("es查询，IO异常！");
        }

        if (RestStatus.OK.equals(searchResponse.status())) {
            // 解析对象
            SearchHit[] hits = searchResponse.getHits().getHits();
            // 获取source
            List<Map<String, Object>> sourceList = Arrays.stream(hits).map(SearchHit::getSourceAsMap).collect(Collectors.toList());
            long totalHits = searchResponse.getHits().getTotalHits().value;
            return new EsQueryRespPO(queryPO.getPageNum(), queryPO.getPageSize(), totalHits, sourceList);
        } else {
            log.error("es查询返回的状态码异常！searchResponse.status={}, index={}, 请求source={}", searchResponse.status(),
                    queryPO.getIndex(), searchRequest.source());
            throw new EsException("es查询返回的状态码异常");
        }
    }

}
